{-# LANGUAGE TupleSections #-}

-- | Branch and Bound with Bernstein polynomials

module BsBB (T, prepare, trm) where

import qualified Bs
import qualified Iv
import qualified BB
import qualified Tree
import qualified Expr
import qualified Ply
import Interlude

data T a = Mk [Bs.T a] [(Bool,Bool)] [(Int,Hand,Bool)] (Iv.T a)

-- | Make a box after carrying out all valid projections, beginning
--   with the dimension specified in the last argument. If the box can
--   be eliminated return Nothing.
box :: (Fractional a, Ord a) => [Bs.T a] -> [(Bool,Bool)] -> [(Int,Hand,Bool)] -> Int -> Maybe (T a)
box bs p@(length->n) pth k = maybe (Just bx) (\(i,h) ->
      if hand fst snd h (p!!i) then Nothing else
      if n == 1 then Just bx else
      box (Bs.proj i h <$> bs') (delIx i p) ((i,h,True):pth) (i `mod` (n-1)))
    $ msum $ (`map` [k..k+n-1]) $ \((`mod` n) -> i) ->
      (i,) <$> (join $ allEq $ map (Bs.sharp i) bs') where
  bx = Mk bs' p pth $ Iv.mk lo (Bs.ublb bs')
  bs' = filter ((lo <=) . Bs.hi) bs -- Disjunct reduction
  lo = foldl1 max (map Bs.lo bs)

instance (Fractional a, Ord a) => BB.C (T a) where
  type BB.Bound (T a) = a

  branch (Mk bs p pth _) =
      catMaybes [ box bs1 (p0++(l,True):p') ((k,L,False):pth) k
                , box bs2 (p0++(True,r):p') ((k,R,False):pth) k ] where
    (p0,(l,r):p') = splitAt k p
    (bs1,bs2) = unzip $ map (Bs.subdiv k) bs
    k = (`maxBy` [0..length p-1]) $
      uncurry (max `on` abs) . Iv.dest . Bs.lohi . (`Bs.diff` minBy Bs.lo bs)

  bounds (Mk _ _ _ bnds) = bnds

prepare :: (Fractional a, Ord a) => Int -> String -> Writer String (T a)
prepare v s = do
  when (v>=1)          $ tellLn "Polynomial mode.\nParsing..."
  foldr seq `flip` ps
     $ when (v>=2)     $ metrics (map Ply.toTree ps)
  when (v>=1)          $ tellLn "Converting..."
  foldr seq `flip` bs
     $ when (v>=2)     $ metrics (map Bs.toTree bs)
  when (v>=1)          $ tellLn "Bounding..."
  return $ fromJust $ box bs (replicate (length x) (False,False)) [] 0
  where (x,es) = justParse (Expr.prob False) s
        ps = map (Expr.evalP (\i -> Ply.var i (x!!i)) Ply.const) es
        bs = map Bs.ofPly ps

metrics :: Show a => [Tree.T a] -> Writer String ()
metrics ts = do
  tellLn $ "        Degree: " ++ unwords (map (\t -> show (map (+ (-1)) (Tree.dim t),
                                            Tree.totalDegree t)) ts)
  tellLn $ "  Coefficients: " ++ unwords (map (show . Tree.countLfs) ts)
  tellLn $ "      Sparsity: " ++ unwords (map (flip (showFFloat (Just 3)) "" . Tree.sparsity) ts)

-- | Termination criterion
trm :: (Ord a, Num a) => Maybe a -> Iv.T a -> Bool
trm Nothing    = (0<) . Iv.lo
trm (Just eps) = (<eps) . Iv.width

instance Show a => Show (T a) where
  show (Mk _ _ pth (Iv.dest -> (a,b))) =
    printf "[%-22s, %-22s] %s" (show a) (show b) $
    foldr (\(i,show->h,p) -> (++ show i ++ bool h (map toLower h) p)) "" pth
